{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# eval the depth of sintel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import os\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '0'\n",
    "from dust3r.depth_eval import depth_evaluation, group_by_directory\n",
    "import numpy as np\n",
    "import cv2\n",
    "from tqdm import tqdm\n",
    "import glob\n",
    "TAG_FLOAT = 202021.25\n",
    "\n",
    "def depth_read(filename):\n",
    "    \"\"\" Read depth data from file, return as numpy array. \"\"\"\n",
    "    f = open(filename,'rb')\n",
    "    check = np.fromfile(f,dtype=np.float32,count=1)[0]\n",
    "    assert check == TAG_FLOAT, ' depth_read:: Wrong tag in flow file (should be: {0}, is: {1}). Big-endian machine? '.format(TAG_FLOAT,check)\n",
    "    width = np.fromfile(f,dtype=np.int32,count=1)[0]\n",
    "    height = np.fromfile(f,dtype=np.int32,count=1)[0]\n",
    "    size = width*height\n",
    "    assert width > 0 and height > 0 and size > 1 and size < 100000000, ' depth_read:: Wrong input size (width = {0}, height = {1}).'.format(width,height)\n",
    "    depth = np.fromfile(f,dtype=np.float32,count=-1).reshape((height,width))\n",
    "    return depth\n",
    "\n",
    "pred_pathes = glob.glob(\"results/sintel_video_depth/*/*.npy\") #TODO: update the path to your prediction\n",
    "pred_pathes = sorted(pred_pathes)\n",
    "print(len(pred_pathes))\n",
    "\n",
    "if len(pred_pathes) > 643:\n",
    "    full = True\n",
    "else:\n",
    "    full = False\n",
    "\n",
    "if full:\n",
    "    depth_pathes = glob.glob(f\"data/sintel/training/depth/*/*.dpt\")\n",
    "    depth_pathes = sorted(depth_pathes)\n",
    "else:\n",
    "    seq_list = [\"alley_2\", \"ambush_4\", \"ambush_5\", \"ambush_6\", \"cave_2\", \"cave_4\", \"market_2\", \n",
    "                        \"market_5\", \"market_6\", \"shaman_3\", \"sleeping_1\", \"sleeping_2\", \"temple_2\", \"temple_3\"]\n",
    "    depth_pathes_folder = [f\"data/sintel/training/depth/{seq}\" for seq in seq_list]\n",
    "    depth_pathes = []\n",
    "    for depth_pathes_folder_i in depth_pathes_folder:\n",
    "        depth_pathes += glob.glob(depth_pathes_folder_i + '/*.dpt')\n",
    "    depth_pathes = sorted(depth_pathes)\n",
    "\n",
    "\n",
    "def get_video_results():\n",
    "    grouped_pred_depth = group_by_directory(pred_pathes)\n",
    "    grouped_gt_depth = group_by_directory(depth_pathes)\n",
    "    gathered_depth_metrics = []\n",
    "\n",
    "    for key in tqdm(grouped_pred_depth.keys()):\n",
    "        pd_pathes = grouped_pred_depth[key]\n",
    "        gt_pathes = grouped_gt_depth[key.replace('_pred_depth', '')]\n",
    "\n",
    "        gt_depth = np.stack([depth_read(gt_path) for gt_path in gt_pathes], axis=0)\n",
    "        pr_depth = np.stack([cv2.resize(np.load(pd_path), (gt_depth.shape[2], gt_depth.shape[1]), interpolation=cv2.INTER_CUBIC)\n",
    "                              for pd_path in pd_pathes], axis=0)\n",
    "        \n",
    "        # for depth eval, set align_with_lad2=False to use median alignment; set align_with_lad2=True to use scale&shift alignment\n",
    "        depth_results, error_map, depth_predict, depth_gt = depth_evaluation(pr_depth, gt_depth, max_depth=70, align_with_lad2=True, use_gpu=True, post_clip_max=70)\n",
    "        gathered_depth_metrics.append(depth_results)\n",
    "\n",
    "    depth_log_path = 'tmp.json'\n",
    "    average_metrics = {\n",
    "        key: np.average(\n",
    "            [metrics[key] for metrics in gathered_depth_metrics], \n",
    "            weights=[metrics['valid_pixels'] for metrics in gathered_depth_metrics]\n",
    "        )\n",
    "        for key in gathered_depth_metrics[0].keys() if key != 'valid_pixels'\n",
    "    }\n",
    "    print('Average depth evaluation metrics:', average_metrics)\n",
    "    \n",
    "get_video_results()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# eval the depth of bonn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import os\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '0'\n",
    "from dust3r.depth_eval import depth_evaluation, group_by_directory\n",
    "import numpy as np\n",
    "import cv2\n",
    "import json\n",
    "from tqdm import tqdm\n",
    "import glob\n",
    "from PIL import Image\n",
    "\n",
    "\n",
    "def depth_read(filename):\n",
    "    # loads depth map D from png file\n",
    "    # and returns it as a numpy array\n",
    "    depth_png = np.asarray(Image.open(filename))\n",
    "    # make sure we have a proper 16bit depth map here.. not 8bit!\n",
    "    assert np.max(depth_png) > 255\n",
    "    depth = depth_png.astype(np.float64) / 5000.0\n",
    "    depth[depth_png == 0] = -1.0\n",
    "    return depth\n",
    "\n",
    "seq_list = [\"balloon2\", \"crowd2\", \"crowd3\", \"person_tracking2\", \"synchronous\"]\n",
    "\n",
    "img_pathes_folder = [f\"data/bonn/rgbd_bonn_dataset/rgbd_bonn_{seq}/rgb_110/*.png\" for seq in seq_list]\n",
    "img_pathes = []\n",
    "for img_pathes_folder_i in img_pathes_folder:\n",
    "    img_pathes += glob.glob(img_pathes_folder_i)\n",
    "img_pathes = sorted(img_pathes)\n",
    "depth_pathes_folder = [f\"data/bonn/rgbd_bonn_dataset/rgbd_bonn_{seq}/depth_110/*.png\" for seq in seq_list]\n",
    "depth_pathes = []\n",
    "for depth_pathes_folder_i in depth_pathes_folder:\n",
    "    depth_pathes += glob.glob(depth_pathes_folder_i)\n",
    "depth_pathes = sorted(depth_pathes)\n",
    "pred_pathes = glob.glob(\"results/bonn_video_depth/*/*.npy\")    #TODO: update the path to your prediction\n",
    "pred_pathes = sorted(pred_pathes)\n",
    "\n",
    "def get_video_results():\n",
    "    grouped_pred_depth = group_by_directory(pred_pathes)\n",
    "    grouped_gt_depth = group_by_directory(depth_pathes, idx=-2)\n",
    "    gathered_depth_metrics = []\n",
    "    print(grouped_gt_depth.keys())\n",
    "    print(grouped_pred_depth.keys())\n",
    "    for key in tqdm(grouped_gt_depth.keys()):\n",
    "        pd_pathes = grouped_pred_depth[key]\n",
    "        gt_pathes = grouped_gt_depth[key]\n",
    "        gt_depth = np.stack([depth_read(gt_path) for gt_path in gt_pathes], axis=0)\n",
    "        pr_depth = np.stack([cv2.resize(np.load(pd_path), (gt_depth.shape[2], gt_depth.shape[1]), interpolation=cv2.INTER_CUBIC)\n",
    "                              for pd_path in pd_pathes], axis=0)\n",
    "        \n",
    "        # for depth eval, set align_with_lad2=False to use median alignment; set align_with_lad2=True to use scale&shift alignment\n",
    "        depth_results, error_map, depth_predict, depth_gt = depth_evaluation(pr_depth, gt_depth, max_depth=70, align_with_lad2=True, use_gpu=True)\n",
    "        gathered_depth_metrics.append(depth_results)\n",
    "\n",
    "    depth_log_path = 'tmp.json'\n",
    "    average_metrics = {\n",
    "        key: np.average(\n",
    "            [metrics[key] for metrics in gathered_depth_metrics], \n",
    "            weights=[metrics['valid_pixels'] for metrics in gathered_depth_metrics]\n",
    "        )\n",
    "        for key in gathered_depth_metrics[0].keys() if key != 'valid_pixels'\n",
    "    }\n",
    "    print('Average depth evaluation metrics:', average_metrics)\n",
    "    with open(depth_log_path, 'w') as f:\n",
    "        f.write(json.dumps(average_metrics))\n",
    "\n",
    "get_video_results()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# eval the depth of kitti"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import os\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '0'\n",
    "from dust3r.depth_eval import depth_evaluation, group_by_directory\n",
    "import numpy as np\n",
    "import cv2\n",
    "import json\n",
    "from tqdm import tqdm\n",
    "import glob\n",
    "from PIL import Image\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "\n",
    "def depth_read(filename):\n",
    "    # loads depth map D from png file\n",
    "    # and returns it as a numpy array,\n",
    "    # for details see readme.txt\n",
    "    img_pil = Image.open(filename)\n",
    "    depth_png = np.array(img_pil, dtype=int)\n",
    "    # make sure we have a proper 16bit depth map here.. not 8bit!\n",
    "    assert(np.max(depth_png) > 255)\n",
    "\n",
    "    depth = depth_png.astype(float) / 256.\n",
    "    depth[depth_png == 0] = -1.\n",
    "    return depth\n",
    "\n",
    "depth_pathes = glob.glob(\"data/kitti/depth_selection/val_selection_cropped/groundtruth_depth_gathered/*/*.png\")\n",
    "depth_pathes = sorted(depth_pathes)\n",
    "pred_pathes = glob.glob(\"results/kitti_video_depth/*/frame_*.npy\") #TODO: update the path to your prediction\n",
    "pred_pathes = sorted(pred_pathes)\n",
    "\n",
    "\n",
    "def get_video_results():\n",
    "    grouped_pred_depth = group_by_directory(pred_pathes)\n",
    "    grouped_gt_depth = group_by_directory(depth_pathes)\n",
    "    gathered_depth_metrics = []\n",
    "    for key in tqdm(grouped_pred_depth.keys()):\n",
    "        pd_pathes = grouped_pred_depth[key]\n",
    "        gt_pathes = grouped_gt_depth[key]\n",
    "        gt_depth = np.stack([depth_read(gt_path) for gt_path in gt_pathes], axis=0)\n",
    "        pr_depth = np.stack([cv2.resize(np.load(pd_path), (gt_depth.shape[2], gt_depth.shape[1]), interpolation=cv2.INTER_CUBIC)\n",
    "                              for pd_path in pd_pathes], axis=0)\n",
    "\n",
    "        # for depth eval, set align_with_lad2=False to use median alignment; set align_with_lad2=True to use scale&shift alignment\n",
    "        depth_results, error_map, depth_predict, depth_gt = depth_evaluation(pr_depth, gt_depth, max_depth=None, align_with_lad2=True, use_gpu=True)\n",
    "        gathered_depth_metrics.append(depth_results)\n",
    "\n",
    "    depth_log_path = 'tmp.json'\n",
    "    average_metrics = {\n",
    "        key: np.average(\n",
    "            [metrics[key] for metrics in gathered_depth_metrics], \n",
    "            weights=[metrics['valid_pixels'] for metrics in gathered_depth_metrics]\n",
    "        )\n",
    "        for key in gathered_depth_metrics[0].keys() if key != 'valid_pixels'\n",
    "    }\n",
    "    print('Average depth evaluation metrics:', average_metrics)\n",
    "    with open(depth_log_path, 'w') as f:\n",
    "        f.write(json.dumps(average_metrics))\n",
    "\n",
    "get_video_results()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# eval the depth of nyu-v2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "from dust3r.depth_eval import depth_evaluation\n",
    "import numpy as np\n",
    "import json\n",
    "from tqdm import tqdm\n",
    "import glob\n",
    "\n",
    "depth_pathes = glob.glob(\"data/nyu-v2/val/nyu_depths/*.npy\")\n",
    "depth_pathes = sorted(depth_pathes)\n",
    "pred_pathes = glob.glob(\"results/nyuv2_depth/*.npy\")   #TODO: update the path to your prediction\n",
    "pred_pathes = sorted(pred_pathes)\n",
    "gathered_depth_metrics = []\n",
    "for idx in tqdm(range(len(depth_pathes))):\n",
    "    pred_depth = np.load(pred_pathes[idx])\n",
    "    gt_depth = np.load(depth_pathes[idx])\n",
    "    pred_depth = cv2.resize(pred_depth, (gt_depth.shape[1], gt_depth.shape[0]), interpolation=cv2.INTER_CUBIC)\n",
    "\n",
    "    depth_results, error_map, depth_predict, depth_gt = depth_evaluation(pred_depth, gt_depth, max_depth=None, lr=1e-3)\n",
    "    gathered_depth_metrics.append(depth_results)\n",
    "\n",
    "depth_log_path = 'tmp.json'\n",
    "average_metrics = {\n",
    "    key: np.average(\n",
    "        [metrics[key] for metrics in gathered_depth_metrics], \n",
    "        weights=[metrics['valid_pixels'] for metrics in gathered_depth_metrics]\n",
    "    )\n",
    "    for key in gathered_depth_metrics[0].keys() if key != 'valid_pixels'\n",
    "}\n",
    "print('Average depth evaluation metrics:', average_metrics)\n",
    "with open(depth_log_path, 'w') as f:\n",
    "    f.write(json.dumps(average_metrics))\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "dust3r",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
